# Obtained from: https://github.com/lhoyer/DAFormer
# ---------------------------------------------------------------
# Copyright (c) 2021-2022 ETH Zurich, Lukas Hoyer. All rights reserved.
# Licensed under the Apache License, Version 2.0
# ---------------------------------------------------------------

import torch
import torch.nn as nn
from mmcv.cnn import ConvModule, DepthwiseSeparableConvModule

from mmseg.models.decode_heads.isa_head import ISALayer
from mmseg.ops import resize
from ..builder import HEADS
from .aspp_head import ASPPModule
from .decode_head import BaseDecodeHead
from .segformer_head import MLP
from .sep_aspp_head import DepthwiseSeparableASPPModule


class ASPPWrapper(nn.Module):


    def __init__(self,
                 in_channels,
                 channels,
                 sep,
                 dilations,
                 pool,
                 norm_cfg,
                 act_cfg,
                 align_corners,
                 context_cfg=None):
        """
            包括：池化层、上下文层、aspp模块、瓶颈层

            参数:
            - in_channels (int): 输入通道数。
            - channels (int): 输出通道数。
            - sep (bool): 是否使用深度可分离卷积。
            - dilations (list|tuple): 空洞卷积的扩张率列表。
            - pool (bool): 是否使用全局平均池化。
            - norm_cfg (dict): 标准化层的配置。
            - act_cfg (dict): 激活函数的配置。
            - align_corners (bool): 是否对齐角落。
            - context_cfg (dict|None): 上下文层的配置，如果为None则不使用上下文层。
        """
        # 初始化父类
        super(ASPPWrapper, self).__init__()

        # 确保扩张率是一个列表或元组
        assert isinstance(dilations, (list, tuple))
        self.dilations = dilations
        self.align_corners = align_corners

        # 根据是否使用全局平均池化来构建相应的层
        if pool:
            self.image_pool = nn.Sequential(
                nn.AdaptiveAvgPool2d(1),
                ConvModule(
                    in_channels,
                    channels,
                    1,
                    norm_cfg=norm_cfg,
                    act_cfg=act_cfg))
        else:
            self.image_pool = None

        # 根据上下文层配置是否存在来决定是否构建上下文层
        if context_cfg is not None:
            self.context_layer = build_layer(in_channels, channels,
                                             **context_cfg)
        else:
            self.context_layer = None

        # 根据是否使用深度可分离卷积来选择ASPP模块的类型
        ASPP = {True: DepthwiseSeparableASPPModule, False: ASPPModule}[sep]
        self.aspp_modules = ASPP(
            dilations=dilations,
            in_channels=in_channels,
            channels=channels,
            norm_cfg=norm_cfg,
            conv_cfg=None,
            act_cfg=act_cfg)

        # 构建瓶颈层，整合所有并行路径的输出
        self.bottleneck = ConvModule(
            (len(dilations) + int(pool) + int(bool(context_cfg))) * channels,
            channels,
            kernel_size=3,
            padding=1,
            norm_cfg=norm_cfg,
            act_cfg=act_cfg)


    def forward(self, x):
        """Forward function."""
        aspp_outs = []
        if self.image_pool is not None:
            aspp_outs.append(
                resize(
                    self.image_pool(x),
                    size=x.size()[2:],
                    mode='bilinear',
                    align_corners=self.align_corners))
        if self.context_layer is not None:
            aspp_outs.append(self.context_layer(x))
        aspp_outs.extend(self.aspp_modules(x))# 使用标准的ASPP模块
        aspp_outs = torch.cat(aspp_outs, dim=1)

        output = self.bottleneck(aspp_outs)# 通过瓶颈层卷积融合多尺度特征
        return output


def build_layer(in_channels, out_channels, type, **kwargs):
    """
    根据指定的类型构建神经网络层。

    参数:
    - in_channels (int): 输入通道数。
    - out_channels (int): 输出通道数。
    - type (str): 网络层的类型，包括'id', 'mlp', 'sep_conv', 'conv', 'aspp', 'rawconv_and_aspp', 'isa'。
    - **kwargs: 其他可变关键字参数，用于传递给特定类型的网络层。

    返回:
    - nn.Module: 构建的网络层模块。
    """
    # 根据不同的类型构建相应的网络层
    if type == 'id':
        # 返回一个恒等层，不做任何变换
        return nn.Identity()
    elif type == 'mlp':
        # 构建一个多层感知器（MLP）层
        return MLP(input_dim=in_channels, embed_dim=out_channels)
    elif type == 'sep_conv':
        # 构建一个深度可分离卷积层
        return DepthwiseSeparableConvModule(
            in_channels=in_channels,
            out_channels=out_channels,
            padding=kwargs['kernel_size'] // 2,
            **kwargs)
    elif type == 'conv':
        # 构建一个常规的卷积层
        return ConvModule(
            in_channels=in_channels,
            out_channels=out_channels,
            padding=kwargs['kernel_size'] // 2,
            **kwargs)
    elif type == 'aspp':
        # 构建一个带有ASPP（空洞空间金字塔池化）的层
        return ASPPWrapper(
            in_channels=in_channels, channels=out_channels, **kwargs)
    elif type == 'rawconv_and_aspp':
        # 先构建一个常规卷积层，后接ASPP层
        kernel_size = kwargs.pop('kernel_size')
        return nn.Sequential(
            nn.Conv2d(
                in_channels=in_channels,
                out_channels=out_channels,
                kernel_size=kernel_size,
                padding=kernel_size // 2),
            ASPPWrapper(
                in_channels=out_channels, channels=out_channels, **kwargs))
    elif type == 'isa':
        # 构建一个ISA（交互感知注意力）层
        return ISALayer(
            in_channels=in_channels, channels=out_channels, **kwargs)
    else:
        # 如果类型不匹配，抛出异常
        raise NotImplementedError(type)



@HEADS.register_module()
class DAFormerHead(BaseDecodeHead):

    def __init__(self, **kwargs):
        """
        初始化DAFormerHead类的构造函数。

        使用可变关键字参数kwargs来增加灵活性，允许在实例化时传递任意的键值对参数。
        """
        # 调用父类的构造方法，传入input_transform参数和所有传入的kwargs参数
        super(DAFormerHead, self).__init__(
            input_transform='multiple_select', **kwargs)

        # 确保align_corners属性为False，以保证一致性
        assert not self.align_corners

        # 从kwargs中提取decoder_params参数，该参数包含了嵌入维度、嵌入配置等信息
        decoder_params = kwargs['decoder_params']
        embed_dims = decoder_params['embed_dims']
        # 如果embed_dims是一个整数，则将其复制为与in_index长度相同的列表
        if isinstance(embed_dims, int):
            embed_dims = [embed_dims] * len(self.in_index)

        # 提取嵌入配置、嵌入颈部配置和融合配置
        embed_cfg = decoder_params['embed_cfg']
        embed_neck_cfg = decoder_params['embed_neck_cfg']
        # 如果embed_neck_cfg设置为与embed_cfg相同，则将其赋值为embed_cfg
        if embed_neck_cfg == 'same_as_embed_cfg':
            embed_neck_cfg = embed_cfg
        fusion_cfg = decoder_params['fusion_cfg']

        # 遍历每个配置文件，如果配置文件不为空且类型为aspp，则设置其align_corners属性
        for cfg in [embed_cfg, embed_neck_cfg, fusion_cfg]:
            if cfg is not None and 'aspp' in cfg['type']:
                cfg['align_corners'] = self.align_corners

        # 初始化嵌入层字典
        self.embed_layers = {}
        # 遍历in_index、in_channels和embed_dims，构建每个嵌入层
        for i, in_channels, embed_dim in zip(self.in_index, self.in_channels,
                                             embed_dims):
            # 如果是最后一个嵌入层，则使用embed_neck_cfg配置
            if i == self.in_index[-1]:
                self.embed_layers[str(i)] = build_layer(
                    in_channels, embed_dim, **embed_neck_cfg)
            else:
                self.embed_layers[str(i)] = build_layer(
                    in_channels, embed_dim, **embed_cfg)
        # 将嵌入层字典转换为ModuleDict，以便后续使用
        self.embed_layers = nn.ModuleDict(self.embed_layers)

        # 构建融合层，将所有嵌入维度的和作为输入通道数，self.channels作为输出通道数
        self.fuse_layer = build_layer(
            sum(embed_dims), self.channels, **fusion_cfg)

    def forward(self, inputs):
        """
        定义前向传播过程，处理输入的多尺度特征图。

        参数:
        - inputs (list[Tensor]): 一个包含多个尺度特征图的列表。

        返回:
        - x (Tensor): 经过融合和分类后的特征图。
        """
        # 初始化输入特征图
        x = inputs
        # 获取最后一个尺度特征图的批量大小、通道数、高度和宽度
        n, _, h, w = x[-1].shape

        # 获取第一个尺度特征图的高度和宽度，用于后续特征图的尺寸调整
        os_size = x[0].size()[2:]
        # 初始化一个字典，用于存储经过嵌入层处理后的特征图
        _c = {}

        # 遍历特征图索引
        for i in self.in_index:
            # 对当前索引的特征图应用相应的嵌入层
            _c[i] = self.embed_layers[str(i)](x[i])
            # 如果处理后的特征图维度为3，对其进行维度变换和重塑
            if _c[i].dim() == 3:
                _c[i] = _c[i].permute(0, 2, 1).contiguous()\
                    .reshape(n, -1, x[i].shape[2], x[i].shape[3])
            # 如果当前特征图的尺寸与第一个特征图不同，对其进行调整
            if _c[i].size()[2:] != os_size:
                _c[i] = resize(
                    _c[i],
                    size=os_size,
                    mode='bilinear',
                    align_corners=self.align_corners)

        # 将所有处理后的特征图在通道维度上进行拼接，并通过融合层进行融合
        x = self.fuse_layer(torch.cat(list(_c.values()), dim=1))
        # 通过分类层生成最终的特征图
        x = self.cls_seg(x)

        # 返回经过分类的特征图
        return x
